
#include "source/aipostprocess.h"
#include "source/logger.h"
#include <math.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>

#include <set>
#include <vector>
//#define LABEL_NALE_TXT_PATH "./coco_80_labels_list.txt"

//static char *labels[OBJ_CLASS_NUM];

inline static int clamp(float val, int min, int max) { return val > min ? (val < max ? val : max) : min; }
char* AiPostProcess::readLine(FILE *fp, char *buffer, int *len) {
    int ch;
    int i = 0;
    size_t buff_len = 0;

    buffer = (char *)malloc(buff_len + 1);
    if (!buffer)
        return NULL; // Out of memory

    while ((ch = fgetc(fp)) != '\n' && ch != EOF) {
        buff_len++;
        void *tmp = realloc(buffer, buff_len + 1);
        if (tmp == NULL) {
            free(buffer);
            return NULL; // Out of memory
        }
        buffer = (char *)tmp;

        buffer[i] = (char)ch;
        i++;
    }
    buffer[i] = '\0';

    *len = buff_len;

    // Detect end
    if (ch == EOF && (i == 0 || ferror(fp))) {
        free(buffer);
        return NULL;
    }
    return buffer;
}

int AiPostProcess::readLines(const char *fileName, char *lines[], int maxLine) {
    FILE *file = fopen(fileName, "r");
    char *s;
    int i = 0;
    int n = 0;

    if (file == NULL) {
        LogError<< "Open file fail: " << fileName;
        return -1;
    }

    while ((s = readLine(file, s, &n)) != NULL) {
        lines[i++] = s;
        if (i >= maxLine)
            break;
    }
    fclose(file);
    return i;
}

int AiPostProcess::loadLabelName(const char *locationFilename, char *label[], int objClassNum) {
    LogInfo << "Load label: " << locationFilename;
    readLines(locationFilename, label, objClassNum);
    return 0;
}



float AiPostProcess::calculateOverlap(float xmin0, float ymin0, float xmax0, float ymax0, float xmin1, float ymin1, float xmax1,
                                      float ymax1) {
    float w = fmax(0.f, fmin(xmax0, xmax1) - fmax(xmin0, xmin1) + 1.0);
    float h = fmax(0.f, fmin(ymax0, ymax1) - fmax(ymin0, ymin1) + 1.0);
    float i = w * h;
    float u = (xmax0 - xmin0 + 1.0) * (ymax0 - ymin0 + 1.0) + (xmax1 - xmin1 + 1.0) * (ymax1 - ymin1 + 1.0) - i;
    return u <= 0.f ? 0.f : (i / u);
}

int AiPostProcess::nms(int validCount, std::vector<float> &outputLocations,
                       std::vector<int> classIds, std::vector<int> &order,
                       int filterId, float threshold)
{
    for (int i = 0; i < validCount; ++i)
    {
        if (order[i] == -1 || classIds[i] != filterId)
        {
            continue;
        }
        int n = order[i];
        for (int j = i + 1; j < validCount; ++j)
        {
            int m = order[j];
            if (m == -1 || classIds[i] != filterId)
            {
                continue;
            }
            float xmin0 = outputLocations[n * 4 + 0];
            float ymin0 = outputLocations[n * 4 + 1];
            float xmax0 = outputLocations[n * 4 + 0] + outputLocations[n * 4 + 2];
            float ymax0 = outputLocations[n * 4 + 1] + outputLocations[n * 4 + 3];

            float xmin1 = outputLocations[m * 4 + 0];
            float ymin1 = outputLocations[m * 4 + 1];
            float xmax1 = outputLocations[m * 4 + 0] + outputLocations[m * 4 + 2];
            float ymax1 = outputLocations[m * 4 + 1] + outputLocations[m * 4 + 3];

            float iou = calculateOverlap(xmin0, ymin0, xmax0, ymax0, xmin1, ymin1, xmax1, ymax1);

            if (iou > threshold)
            {
                order[j] = -1;
            }
        }
    }
    return 0;
}

int AiPostProcess::quickSortIndiceInverse(std::vector<float> &input, int left, int right, std::vector<int> &indices)
{
    float key;
    int key_index;
    int low = left;
    int high = right;
    if (left < right)
    {
        key_index = indices[left];
        key = input[left];
        while (low < high)
        {
            while (low < high && input[high] <= key)
            {
                high--;
            }
            input[low] = input[high];
            indices[low] = indices[high];
            while (low < high && input[low] >= key)
            {
                low++;
            }
            input[high] = input[low];
            indices[high] = indices[low];
        }
        input[low] = key;
        indices[low] = key_index;
        quickSortIndiceInverse(input, left, low - 1, indices);
        quickSortIndiceInverse(input, low + 1, right, indices);
    }
    return low;
}

float AiPostProcess::sigmoid(float x) {
    return 1.0 / (1.0 + expf(-x));
}

float AiPostProcess::unsigmoid(float y) {
    return -1.0 * logf((1.0 / y) - 1.0);
}

int32_t AiPostProcess::clip(float val, float min, float max) {
    float f = val <= min ? min : (val >= max ? max : val);
    return f;
}

int8_t AiPostProcess::qntF32ToAffine(float f32, int32_t zp, float scale) {
    float dst_val = (f32 / scale) + zp;
    int8_t res = (int8_t)clip(dst_val, -128, 127);
    return res;
}

float AiPostProcess::deqntAffineToF32(int8_t qnt, int32_t zp, float scale) {
    return ((float)qnt - (float)zp) * scale;
}


void AiPostProcess::computeDfl(float* tensor, int dflLen, float* box){
    for (int b=0; b<4; b++){
        float exp_t[dflLen];
        float exp_sum=0;
        float acc_sum=0;
        for (int i=0; i< dflLen; i++){
            exp_t[i] = exp(tensor[i+b*dflLen]);
            exp_sum += exp_t[i];
        }

        for (int i=0; i< dflLen; i++){
            acc_sum += exp_t[i]/exp_sum *i;
        }
        box[b] = acc_sum;
    }
}

// int8_t *box_tensor: 边界框数据的量化张量。
// int32_t box_zp: 边界框量化张量的零点（Zero Point）。
// float box_scale: 边界框量化张量的尺度（Scale）。
// int8_t *score_tensor: 置信度数据的量化张量。
// int32_t score_zp: 置信度量化张量的零点。
// float score_scale: 置信度量化张量的尺度。
// int8_t *score_sum_tensor: 用于快速过滤的置信度和量化张量。
// int32_t score_sum_zp: 置信度和量化张量的零点。
// float score_sum_scale: 置信度和量化张量的尺度。
// int grid_h: 网格的高度。
// int grid_w: 网格的宽度。
// int stride: 步长，表示特征图到输入图像的缩放比例。
// int dfl_len: 边界框解码所需的 DFL（Distribution Focal Loss）长度。
// std::vector<float> &boxes: 存储解码后的边界框。
// std::vector<float> &objProbs: 存储目标置信度。
// std::vector<int> &classId: 存储类别 ID。
// float threshold: 置信度阈值，用于过滤低置信度的检测框。
int AiPostProcess::processI8(int8_t *boxTensor, int32_t boxZp, float boxScale, int8_t *scoreTensor, int32_t scoreZp, float scoreScale,
                             int8_t *scoreSumTensor, int32_t scoreSumZp, float scoreSumScale, int gridH, int gridW, int stride, int dflLen,
                             std::vector<float> &boxes, std::vector<float> &objProbs, std::vector<int> &classId, float threshold)
{
    int validCount = 0;
    int gridLen = gridH * gridW;
    int8_t scoreThresI8 = qntF32ToAffine(threshold, scoreZp, scoreScale);
    int8_t scoreSumThresI8 = qntF32ToAffine(threshold, scoreSumZp, scoreSumScale);
    // LogDebug << "scoreThresI8 :" << (int)scoreThresI8;
    // LogDebug << "scoreSumThresI8 :" << (int)scoreSumThresI8;
    // LogDebug << "scoreZp:" <<( int)scoreZp;
    // LogDebug << "scoreSumZp:" << (int)scoreSumZp;

    for (int i = 0; i < gridH; i++)  // 遍历每个网格的高度
    {
        for (int j = 0; j < gridW; j++)  // 遍历每个网格的宽度
        {
            int offset = i * gridW + j;  // 当前网格的偏移量
            int max_class_id = -1;  // 存储最大置信度对应的类别 ID

            // 通过 score sum 起到快速过滤的作用
            if (scoreSumTensor  != nullptr)
            {
                if (scoreSumTensor[offset] < scoreSumThresI8)
                {
                    continue;  // 如果当前网格的置信度和低于阈值，则跳过该网格
                }
            }

            int8_t maxScore = -scoreZp;  // 初始化最大置信度为最小值

            // qDebug() << "objClassNum:" << objClassNum;
            for (int c = 0; c < objClassNum; c++)  // 遍历所有类别
            {
                // qDebug() << "c" << c << "objClassNum:" << objClassNum;
                // LogDebug << "scoreTensor[offset] : " << (int)scoreTensor[offset]
                //          << "scoreThresI8 : "  << (int)scoreThresI8
                //          <<" maxScore" << (int)maxScore
                //          <<" offset" << (int)offset;
                // qDebug() << "offset:" << offset << "scoreTensor.size:" << sizeof(scoreTensor) ;
                if ((scoreTensor[offset] > scoreThresI8) && (scoreTensor[offset] > maxScore))
                {
                    maxScore = scoreTensor[offset];  // 更新最大置信度
                    max_class_id = c;  // 更新最大置信度对应的类别 ID
                }
                offset += gridLen;  // 移动到下一个类别的偏移量
            }

            // LogDebug <<" maxScore: " << (int)maxScore
            //          << "scoreSumTensor[offset]" <<(int)scoreSumTensor[offset]
            //          <<" scoreThresI8: " << (int)scoreThresI8 << i << "||" << j;

            // 计算边界框
            if (maxScore > scoreThresI8)
            {
                //LogDebug << "offset" << offset;
                offset = i * gridW + j;  // 重新计算偏移量
                float box[4];
                float before_dfl[dflLen * 4];
                //LogDebug << "goin4";
                for (int k = 0; k < dflLen * 4; k++)
                {
                    before_dfl[k] = deqntAffineToF32(boxTensor[offset], boxZp, boxScale);  // 解量化边界框数据
                    offset += gridLen;  // 移动到下一个数据点
                }
                //LogDebug << "goin5";
                for(int a=0 ; a< 5;a++)
                    //LogDebug << "before_dfl" << a << ':' << before_dfl[a];

                computeDfl(before_dfl, dflLen, box);  // 计算边界框位置
                //LogDebug << "goin2";
                float x1, y1, x2, y2, w, h;
                x1 = (-box[0] + j + 0.5) * stride;
                y1 = (-box[1] + i + 0.5) * stride;
                x2 = (box[2] + j + 0.5) * stride;
                y2 = (box[3] + i + 0.5) * stride;
                w = x2 - x1;
                h = y2 - y1;
                boxes.push_back(x1);  // 存储边界框的左上角 x 坐标
                boxes.push_back(y1);  // 存储边界框的左上角 y 坐标
                boxes.push_back(w);  // 存储边界框的宽度
                boxes.push_back(h);  // 存储边界框的高度
                // LogDebug << "goin3";
                objProbs.push_back(deqntAffineToF32(maxScore, scoreZp, scoreScale));  // 存储置信度
                classId.push_back(max_class_id);  // 存储类别 ID
                validCount++;  // 增加有效检测框的计数
            }
        }
    }
    //LogDebug << "validCount:" << validCount;
    return validCount;  // 返回有效检测框的数量
}

int AiPostProcess::processFp32(float *boxTensor, float *scoreTensor, float *scoreSumTensor, int gridH, int gridW, int stride, int dflLen,
                               std::vector<float> &boxes, std::vector<float> &objProbs, std::vector<int> &classId, float threshold) {
    int validCount = 0;
    int gridLen = gridH * gridW;
    for (int i = 0; i < gridH; i++) {
        for (int j = 0; j < gridW; j++) {
            int offset = i * gridW + j;
            int maxClassId = -1;

            if (scoreSumTensor != nullptr) {
                if (scoreSumTensor[offset] < threshold) {
                    continue;
                }
            }

            float maxScore = 0;
            for (int c = 0; c < objClassNum; c++) {
                if ((scoreTensor[offset] > threshold) && (scoreTensor[offset] > maxScore)) {
                    maxScore = scoreTensor[offset];
                    maxClassId = c;
                }
                offset += gridLen;
            }

            if (maxScore > threshold) {
                offset = i * gridW + j;
                float box[4];
                float beforeDfl[dflLen * 4];
                for (int k = 0; k < dflLen * 4; k++) {
                    beforeDfl[k] = boxTensor[offset];
                    offset += gridLen;
                }
                computeDfl(beforeDfl, dflLen, box);

                float x1, y1, x2, y2, w, h;
                x1 = (-box[0] + j + 0.5) * stride;
                y1 = (-box[1] + i + 0.5) * stride;
                x2 = (box[2] + j + 0.5) * stride;
                y2 = (box[3] + i + 0.5) * stride;
                w = x2 - x1;
                h = y2 - y1;
                boxes.push_back(x1);
                boxes.push_back(y1);
                boxes.push_back(w);
                boxes.push_back(h);

                objProbs.push_back(maxScore);
                classId.push_back(maxClassId);
                validCount++;
            }
        }
    }
    return validCount;
}

int AiPostProcess::postProcess(rknn_app_context_t *appCtx, rknn_tensor_mem **outputs, letterbox_t *letterBox, float confThreshold, float nmsThreshold, object_detect_result_list *odResults) {
    LogTrack << "in postProcess";
    std::vector<float> filterBoxes;
    std::vector<float> objProbs;
    std::vector<int> classId;
    int validCount = 0;
    int stride = 0;
    int gridH = 0;
    int gridW = 0;
    int modelInW = appCtx->model_width;
    int modelInH = appCtx->model_height;

    memset(odResults, 0, sizeof(object_detect_result_list));

    int dflLen = appCtx->output_attrs[0].dims[1] / 4;


    int outputPerBranch = appCtx->io_num.n_output / 3;
    // qDebug() << "appCtx->io_num.n_output :" << appCtx->io_num.n_output;
    for (int i = 0; i < outputPerBranch; i++) {
        void *scoreSum = nullptr;
        int32_t scoreSumZp = 0;
        float scoreSumScale = 1.0;
        // if (outputPerBranch == 3) {
            scoreSum = outputs[i * 3 + 2]->virt_addr;
            scoreSumZp = appCtx->output_attrs[i * 3 + 2].zp;
            scoreSumScale = appCtx->output_attrs[i * 3 + 2].scale;
        // }
        int boxIdx = i * 3;
        int scoreIdx = i * 3 + 1;

        gridH = appCtx->output_attrs[boxIdx].dims[2];
        gridW = appCtx->output_attrs[boxIdx].dims[3];
        stride = modelInH / gridH;
        // qDebug() << "gridH:" << gridH << "gridW:" << gridW;
        if (appCtx->is_quant) {
            validCount += processI8((int8_t *)outputs[boxIdx]->virt_addr, appCtx->output_attrs[boxIdx].zp, appCtx->output_attrs[boxIdx].scale,
                                    (int8_t *)outputs[scoreIdx]->virt_addr, appCtx->output_attrs[scoreIdx].zp, appCtx->output_attrs[scoreIdx].scale,
                                    (int8_t *)scoreSum, scoreSumZp, scoreSumScale, gridH, gridW, stride, dflLen,
                                    filterBoxes, objProbs, classId, confThreshold);
        } else {
            validCount += processFp32((float *)outputs[boxIdx]->virt_addr, (float *)outputs[scoreIdx]->virt_addr, (float *)scoreSum,
                                      gridH, gridW, stride, dflLen, filterBoxes, objProbs, classId, confThreshold);
        }
    }
    // for(int a=0 ; a < 20;a++)
    // {
    //     LogDebug << "indexArrachar *y : " << filterBoxes[a];
    // }
    LogTrack << "validCount:validCount: " << validCount << "QU:" << appCtx->is_quant;
    if (validCount <= 0) {
        return 0;
    }
    std::vector<int> indexArray;
    for (int i = 0; i < validCount; ++i) {
        indexArray.push_back(i);
    }
    quickSortIndiceInverse(objProbs, 0, validCount - 1, indexArray);

    std::set<int> class_set(std::begin(classId), std::end(classId));
    LogTrack << "nmsThreshold: "  << nmsThreshold;


    for (auto c : class_set) {
        nms(validCount, filterBoxes, classId, indexArray, c, nmsThreshold);
    }

    int last_count = 0;
    odResults->count = 0;



    // for(int a=0 ; a < validCount;a++)
    // {
    //     LogDebug << "indexArray : " << indexArray[a];
    // }

    for (int i = 0; i < validCount; ++i) {
        if (indexArray[i] == -1 || last_count >= objNumbMaxSize) {
            continue;
        }
        int n = indexArray[i];

        float x1 = filterBoxes[n * 4 + 0] - letterBox->x_pad;
        float y1 = filterBoxes[n * 4 + 1] - letterBox->y_pad;
        float x2 = x1 + filterBoxes[n * 4 + 2];
        float y2 = y1 + filterBoxes[n * 4 + 3];
        int id = classId[n];
        float obj_conf = objProbs[i];

        odResults->results[last_count].box.left = (int)(clip(x1, 0, modelInW) / letterBox->scale);
        odResults->results[last_count].box.top = (int)(clip(y1, 0, modelInH) / letterBox->scale);
        odResults->results[last_count].box.right = (int)(clip(x2, 0, modelInW) / letterBox->scale);
        odResults->results[last_count].box.bottom = (int)(clip(y2, 0, modelInH) / letterBox->scale);
        odResults->results[last_count].prop = obj_conf;
        odResults->results[last_count].cls_id = id;
        last_count++;
    }
    odResults->count = last_count;
    return 0;
}

AiPostProcess::AiPostProcess(int objNumbMaxSize_, float nmsThresh_, float boxThresh_, int objClassNum_)
    : objClassNum(objClassNum_), objNumbMaxSize(objNumbMaxSize_),
    nmsThresh(nmsThresh_),boxThresh(boxThresh_)
{
    LogDebug << "Init objClassNum" << objClassNum;
}

int AiPostProcess::initPostProcess(const std::string &labelNameTxTPath ,int objClassNum_) {
    int ret = 0;
    ret = loadLabelName(labelNameTxTPath.c_str(), labels, objClassNum_);
    if (ret < 0) {
        LogDebug << "Load failed: " << labelNameTxTPath;
        return -1;
    }
    objClassNum = objClassNum_;
    qInfo() << "objClassNum_:" << objClassNum;
    return 0;
}

char* AiPostProcess::cocoClsToName(int cls_id) {
    //LogDebug << "cls_id:" << cls_id << "objClassNum" << objClassNum;
    if (cls_id >= objClassNum) {
        return "null";
    }
    if (labels[cls_id]) {
        return labels[cls_id];
    }
    return "null";
}

void AiPostProcess::deinitPostProcess() {
    for (int i = 0; i < objClassNum; i++) {
        if (labels[i] != nullptr) {
            free(labels[i]);
            labels[i] = nullptr;
        }
    }
}
